#include "dat_demuxer.h"
#include "media/file_utils.h"
#include "media/log.h"
#include <cstring>
#include <algorithm>

namespace OHOS {
namespace Media {
namespace Demuxer {

DatDemuxer::~DatDemuxer() {
    Release();
}

int32_t DatDemuxer::Init(const std::string& filePath) {
    if (isInited_) {
        MEDIA_LOGW("DatDemuxer already inited");
        return 0;
    }

    // 打开文件
    fileHandle_ = FileUtils::OpenFile(filePath, "rb");
    if (fileHandle_ == -1) {
        MEDIA_LOGE("Open DAT file failed: %s", filePath.c_str());
        return -1;
    }

    // 获取文件大小
    fileSize_ = FileUtils::GetFileSize(fileHandle_);
    if (fileSize_ <= 0) {
        MEDIA_LOGE("Invalid DAT file size");
        Release();
        return -1;
    }

    // 初始化轨道信息(MPEG-PS标准轨道)
    DatTrack videoTrack;
    videoTrack.trackId = 0;
    videoTrack.codecMime = "video/mpeg"; // MPEG-1/2视频
    videoTrack.width = 720;
    videoTrack.height = 480;
    videoTrack.streamId = 0xE0; // 视频流ID
    tracks_.push_back(videoTrack);

    DatTrack audioTrack;
    audioTrack.trackId = 1;
    audioTrack.codecMime = "audio/mpeg"; // MPEG-1音频
    audioTrack.sampleRate = 44100;
    audioTrack.channelCount = 2;
    audioTrack.streamId = 0xC0; // 音频流ID
    tracks_.push_back(audioTrack);

    // 提取字幕轨道
    ExtractSubtitleTrack();

    currentOffset_ = 0;
    isInited_ = true;
    MEDIA_LOGI("DatDemuxer init success, track count: %zu", tracks_.size());
    return 0;
}

int32_t DatDemuxer::GetTrackCount(uint32_t& count) {
    if (!isInited_) {
        MEDIA_LOGE("DatDemuxer not inited");
        return -1;
    }
    count = tracks_.size();
    return 0;
}

int32_t DatDemuxer::GetTrackInfo(uint32_t trackId, TrackInfo& info) {
    if (!isInited_ || trackId >= tracks_.size()) {
        MEDIA_LOGE("Invalid track ID: %u", trackId);
        return -1;
    }

    const auto& track = tracks_[trackId];
    info.trackId = track.trackId;
    info.codecMime = track.codecMime;
    info.width = track.width;
    info.height = track.height;
    info.sampleRate = track.sampleRate;
    info.channelCount = track.channelCount;
    return 0;
}

int32_t DatDemuxer::ReadFrame(uint32_t trackId, MediaPacket& packet) {
    if (!isInited_ || trackId >= tracks_.size()) {
        MEDIA_LOGE("Invalid track ID: %u", trackId);
        return -1;
    }

    const uint8_t targetStreamId = tracks_[trackId].streamId;
    uint8_t buffer[4096];
    size_t totalRead = 0;

    // 循环读取直到找到目标轨道的PES包
    while (currentOffset_ + 4096 < fileSize_) {
        ssize_t readSize = FileUtils::ReadFile(
            fileHandle_, buffer, 4096, currentOffset_);
        
        if (readSize <= 0) {
            return -1; // 文件结束
        }

        size_t headerOffset;
        if (FindPsHeader(buffer, readSize, headerOffset)) {
            // 解析PS包
            if (ParsePsPacket(buffer + headerOffset, readSize - headerOffset)) {
                // 解析PES包
                MediaPacket tempPacket;
                uint8_t streamId;
                if (ParsePesPacket(buffer + headerOffset + 4, readSize - headerOffset - 4, tempPacket, streamId)) {
                    if (streamId == targetStreamId) {
                        // 找到目标轨道的包
                        packet = tempPacket;
                        packet.trackId = trackId;
                        currentOffset_ += headerOffset + 4 + tempPacket.size;
                        return 0;
                    } else {
                        // 释放非目标轨道的包
                        delete[] tempPacket.data;
                    }
                }
            }
        }

        currentOffset_ += readSize;
        totalRead += readSize;
        if (totalRead > 1024 * 1024) { // 防止无限循环
            MEDIA_LOGE("Cannot find target track data");
            return -1;
        }
    }

    return -1;
}

int32_t DatDemuxer::Seek(int64_t timeMs) {
    if (!isInited_) {
        MEDIA_LOGE("DatDemuxer not inited");
        return -1;
    }
    // 简单实现：按时间比例计算偏移
    double ratio = static_cast<double>(timeMs) / 3600000; // 假设1小时时长
    currentOffset_ = static_cast<int64_t>(fileSize_ * ratio);
    MEDIA_LOGI("DAT seek to %lld ms, offset: %lld", (long long)timeMs, (long long)currentOffset_);
    return 0;
}

int32_t DatDemuxer::Release() {
    if (fileHandle_ != -1) {
        FileUtils::CloseFile(fileHandle_);
        fileHandle_ = -1;
    }
    tracks_.clear();
    fileSize_ = 0;
    currentOffset_ = 0;
    isInited_ = false;
    return 0;
}

bool DatDemuxer::FindPsHeader(uint8_t* buffer, size_t bufferSize, size_t& offset) {
    // 查找PS包起始标记(0x000001BA)
    for (offset = 0; offset < bufferSize - 3; offset++) {
        if (buffer[offset] == 0x00 && buffer[offset+1] == 0x00 && 
            buffer[offset+2] == 0x01 && buffer[offset+3] == 0xBA) {
            return true;
        }
    }
    return false;
}

bool DatDemuxer::ParsePsPacket(uint8_t* data, size_t size) {
    if (size < 4) return false;
    // PS包头部最小4字节，这里简单验证
    return true;
}

bool DatDemuxer::ParsePesPacket(uint8_t* data, size_t size, MediaPacket& packet, uint8_t& streamId) {
    if (size < 6) return false;

    // 解析流ID(第1字节)
    streamId = data[0];
    
    // 解析PES包长度(2-3字节)
    uint16_t pesLength = (data[1] << 8) | data[2];
    if (pesLength + 6 > size) {
        MEDIA_LOGE("Invalid PES length");
        return false;
    }

    // 解析PES头部(4-5字节)
    uint8_t pesHeaderLength = data[5];
    if (6 + pesHeaderLength > size) {
        MEDIA_LOGE("Invalid PES header length");
        return false;
    }

    // 解析PTS(如果存在)
    uint32_t pts = 0;
    if ((data[4] & 0xC0) == 0x80) { // 包含PTS
        pts = ((data[6] & 0x0E) << 27) | 
              ((data[7] & 0xFF) << 20) |
              ((data[8] & 0x0E) << 15) |
              ((data[9] & 0xFF) << 8) |
              ((data[10] & 0x0E) << 3);
    }

    // 提取负载数据
    size_t payloadSize = pesLength - pesHeaderLength;
    uint8_t* payload = new uint8_t[payloadSize];
    memcpy(payload, data + 6 + pesHeaderLength, payloadSize);

    // 填充包信息
    packet.data = payload;
    packet.size = payloadSize;
    packet.pts = ConvertPts(pts);
    packet.isKeyFrame = (streamId == 0xE0); // 视频轨默认为关键帧
    return true;
}

bool DatDemuxer::ExtractSubtitleTrack() {
    DatTrack subtitleTrack;
    subtitleTrack.trackId = 2;
    subtitleTrack.codecMime = "text/srt"; // 转换为SRT格式
    subtitleTrack.streamId = 0xBD; // 私有流ID(字幕通常在此)
    tracks_.push_back(subtitleTrack);
    MEDIA_LOGI("Subtitle track extracted");
    return true;
}

int64_t DatDemuxer::ConvertPts(uint32_t pts) {
    // MPEG-PS时间戳基于90kHz，转换为ms
    return (pts * 1000) / 90000;
}

} // namespace Demuxer
} // namespace Media
} // namespace OHOS